import os

import pymysql
import pandas as pd
from dotenv import load_dotenv
# 当做一个外部函数
from zhipuai import ZhipuAI

load_dotenv()
api_key = os.getenv("api_key_zp")


# 连接数据库
def connect_database(query):
    print("sql")
    print(query)
    print("sql")
    # mysql的连接配置
    conn = pymysql.connect(
        host='localhost',
        port=3306,
        user='root',
        password='123456',
        database='school',  # 需要连接的数据库
        charset='utf8mb4',  # 编码格式
    )
    cursor = conn.cursor()  # 链接数据库
    cursor.execute(query)  # 执行你的sql语句
    data = cursor.fetchall()  # 获取查询结果
    cols = cursor.description  # 获取列(column)
    conn.close()  # 关闭链接
    col = []  # 存储的列名
    for i in cols:
        col.append(i[0])
    data = list(map(list, data))
    data = pd.DataFrame(data, columns=col)
    return data


# 使用的是智谱清言
client = ZhipuAI(api_key=api_key)


def get_completion(user_prompt, model="glm-4-plus"):
    instruction = """
    你是一位专业的SQL编写工程师,可以根据表结构和用户输入，生成SQL语句。
    """
    table_structures = """
    表结构如下：
    -- 创建班级表
    CREATE TABLE IF NOT EXISTS class (
        class_id INT AUTO_INCREMENT PRIMARY KEY,
        class_name VARCHAR(50) NOT NULL,
        teacher_name VARCHAR(50) NOT NULL
    ) ENGINE=INNODB DEFAULT CHARSET=utf8;

    -- 创建学生表
    CREATE TABLE IF NOT EXISTS student (
        student_id INT AUTO_INCREMENT PRIMARY KEY,
        student_name VARCHAR(50) NOT NULL,
        gender ENUM('男', '女') NOT NULL,
        age INT NOT NULL,
        class_id INT,
        FOREIGN KEY (class_id) REFERENCES class(class_id)
    ) ENGINE=INNODB DEFAULT CHARSET=utf8;

    -- 创建考试分数表
    CREATE TABLE IF NOT EXISTS exam_score (
        score_id INT AUTO_INCREMENT PRIMARY KEY,
        student_id INT,
        SUBJECT VARCHAR(50) NOT NULL,
        score DECIMAL(5, 2) NOT NULL, -- 假设分数为最多3位整数和2位小数
        exam_date DATE NOT NULL,
        FOREIGN KEY (student_id) REFERENCES student(student_id)
    ) ENGINE=INNODB DEFAULT CHARSET=utf8;
    """
    examples = """
        Q: 用户输入样例：
        A: 获取每个学生的姓名、性别、年龄、班级名称、教师名称、科目、分数和考试日期。

        Q: 生成的SQL样例：
        A: SELECT 
            s.student_id,
            s.student_name,
            s.gender,
            s.age,
            c.class_name,
            c.teacher_name,
            es.SUBJECT,
            es.score,
            es.exam_date
        FROM 
            student s
        INNER JOIN class c ON s.class_id = c.class_id
        INNER JOIN exam_score es ON s.student_id = es.student_id;
    """
    prompt = f"""
        {instruction}
        示例：
        {examples}
        表结构如下：
        {table_structures}
        用户输入：
        {user_prompt}
    """
    # print(prompt)
    messages = [{"role": "user", "content": prompt}]
    response = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=0,  # 模型输出的随机性，0 表示随机性最小
    )
    return response.choices[0].message.content


import re

# 初始化全局变量data
data = []


# 定义提交函数
def submit_fn(user_input):
    print(user_input)
    global data
    # 将用户输入的提示词传入给AI大模型,让他来生成SQL语句
    sql = get_completion(user_input)

    # 我们现在只讲了prompt,AI大模型生成SQL有其他的内容,利用正则表达式,提纯SQL
    # 使用正则表达式提取```内的SQL语句
    match = re.search(r'```sql\n(.*?)\n```', sql, re.DOTALL)

    if match:
        sql = match.group(1)
        print("提取的SQL语句:")
    else:
        print("未找到SQL语句")

    # 从数据库中查询结果
    data = connect_database(sql)

    # 返回更新后的数据
    return data


import gradio as gr

# 创建Gradio Blocks应用
with gr.Blocks() as demo:
    gr.HTML("<h1 align='center'>渲染页面</h1>")
    user_input = gr.Textbox(label="查询的内容")
    # 渲染数据，value表示默认数据
    dataframe = gr.Dataframe(value=pd.DataFrame(data))
    # 当用户在user_input中输入内容时，将内容添加到data中，然后重新渲染
    user_input.submit(fn=submit_fn, inputs=user_input, outputs=dataframe)

if __name__ == '__main__':
    demo.launch()
